#对真值进行统计 并保存为一个csv文件
import json
import pandas as pd
import csv
import numpy as np
from transforms3d.euler import  mat2euler, euler2quat,euler2mat

def se3_mat2euler(pose):
    R=mat2euler(pose[:3,:3],axes='szyx')
    return np.rad2deg(R)

obj_id = 1
split = "test"

for obj_id in [1,2,3]:
    for split in ["test","train_pbr"]:
        scene_gt_json = f"E:/pose/datasets/obj_ac/obj_00000{obj_id}/{split}/000001/scene_gt.json"
        scene_gt = json.load(open(scene_gt_json))
        df = pd.DataFrame(columns=['x', 'y', 'z', 'r'])

        for idx in scene_gt.keys():
            scene_gt[idx] = scene_gt[idx][0]
            cam_R_m2c = np.array(scene_gt[idx]["cam_R_m2c"]).reshape(3,3)
            cam_t_m2c=scene_gt[idx]["cam_t_m2c"]
            x,y,z = se3_mat2euler(cam_R_m2c)
            r = np.linalg.norm(cam_t_m2c)
            df = df.append({'x': x, 'y': y, 'z': z, 'r': r}, ignore_index=True)

        df.to_csv(f'plot/output_{obj_id}_{split}.csv',index_label="idx")

        # Define the bins for the angles
        bins = range(0, 180, 1)  # Change the step size to adjust the granularity of the bins
        counts_df = pd.DataFrame()
        from matplotlib import pyplot as plt
        plt.figure()
        for angle in ['x', 'y', 'z']:
            # Create a new column for the binned angles
            df['binned'] = pd.cut(df[angle], bins)
            # Count the number of occurrences in each bin
            counts = df['binned'].value_counts().sort_index()
            # Print the counts
            counts_df[angle] =counts
            plt.hist(df[angle], bins=180, label=angle)
            plt.legend()
  
        counts_df.to_csv(f'plot/counts_{obj_id}_{split}_hist.csv')

        # # Define the bins for the distances
        # bins = range(0, int(df['r'].max()) + 1, 1)  # Change the step size to adjust the granularity of the bins

        # # Create a new column for the binned distances
        # df['binned'] = pd.cut(df['r'], bins)

        # # Count the number of occurrences in each bin
        # counts = df['binned'].value_counts().sort_index()

        # # Save the counts to a CSV file
        # counts.to_csv(f'plot/counts_{obj_id}_{split}_r.csv')